Overview

This PR replaces per-rank full-model materialization with rank-0-authoritative, FSDP2-sharded loading. Rank 0 loads real CPU weights; other ranks build meta shells, allocate only their local shards, and receive weights from rank 0. This removes the old full-model H2D copy on every rank, substantially reducing initialization latency and peak GPU memory while preserving bitwise-equivalent base-weight shards.

End-to-end loading flow

1. Load the scheduler independently, then process each trainable component separately.
2. Rank 0 uses diffusers' low_cpu_mem_usage=True path and holds the authoritative real CPU weights. Other ranks run under init_empty_weights(include_buffers=False) with low_cpu_mem_usage=False, so parameters remain on meta; their checkpoint state dict may still be read transiently on CPU. LoRA is initialized only on rank 0, while other ranks build meta adapters.
3. Verify that every non-rank-0 parameter is meta, then broadcast rank 0's per-tensor parameter and buffer dtypes so all ranks have an identical sharding layout, including fp32-pinned modules.
4. Capture rank 0's full state, apply FSDP2, move rank 0's real local shards to CUDA, and use to_empty(cuda) to allocate only the other ranks' local shards. set_model_state_dict(..., broadcast_from_rank0=True) fills those shards.
5. Broadcast every named buffer from rank 0, then run model-specific post-materialization hooks. With CPU offload, parameters return to CPU while buffers remain on CUDA for forward.

Important edge cases and fixes

• Meta init vs. low_cpu_mem_usage: non-rank-0 must use False; diffusers' True path materializes parameters by assignment and bypasses the ambient meta context.
• _keep_in_fp32_modules: diffusers rejects low_cpu_mem_usage=False for classes with fp32-pinned modules. The pin is temporarily cleared only on non-rank-0; rank 0 still honors it, and sync_model_dtypes restores the authoritative dtype for every meta parameter before sharding.
• Buffers: include_buffers=False lets constructors compute real buffers, but to_empty discards their values and non-persistent buffers are absent from state_dict. Broadcasting all buffers after parameter loading restores correct, replicated values on every rank.
• Wan optional components: passing transformer=None or transformer_2=None through DiffusionPipeline.from_pretrained does not reliably skip them because Wan declares them as optional defaults; the None entries are dropped and both 14B experts can be loaded unintentionally, including during scheduler loading. The fix resolves the requested class from model_index.json and calls cls.from_pretrained(..., subfolder=component) directly, so only the named component is touched and the fp32 workaround targets the correct class.

Benchmark: meta-init + rank0-broadcast loading vs main (8x H200, cold page cache)

Measured on the real FSDPTrainRayActor.init model-loading path;
checkpoint page cache evicted (posix_fadvise DONTNEED) before every run;
latency and peak memory are the max across ranks.
Per-rank shard sha256 digests are bitwise-identical between the two
modes for all four families (same world size, same sharding plan, so the
local shards must match exactly).

Init latency (s) — lower is better

xychart-beta
    title "init latency (s): main (back) vs stream (front)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "seconds" 0 --> 434
    bar [37.2, 247.6, 377.0, 68.2]
    bar [15.8, 120.2, 133.2, 56.3]

Peak GPU memory per rank (GB) — lower is better

xychart-beta
    title "peak GPU/rank (GB): main (back) vs stream (front)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "GB" 0 --> 88
    bar [9.3, 76.3, 60.1, 35.5]
    bar [2.9, 10.9, 14.6, 5.2]

xychart-beta
    title "init speedup (x, higher is better)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "speedup (x)" 0 --> 3
    bar [2.4, 2.1, 2.8, 1.2]

xychart-beta
    title "peak GPU/rank reduction (%)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "% reduction" 0 --> 99
    bar [68.8, 85.7, 75.7, 85.4]

Loading mechanics: low_cpu_mem_usage, meta init, and _keep_in_fp32_modules

This is the most confusing corner of the PR, so spelling it out.

diffusers' from_pretrained has two loading paths:

• low_cpu_mem_usage=True (diffusers' default): self-contained fast path.
  diffusers builds the module on its own internal meta device, then reads the
  checkpoint and materializes every param directly by assignment. It always
  produces real weights and bypasses any ambient init context (assignment does
  not go through parameter registration, so an external init_empty_weights
  cannot intercept it). Peak CPU ~= 1x model size.
• low_cpu_mem_usage=False: plain path. Runs the model's ordinary
  __init__ (params allocated through normal registration), then reads the
  full state dict and copy_()s it into the params. In normal use this peaks
  at ~2x model size -- that is exactly what True was invented to avoid.

Why each rank uses what it uses:

Non-rank0 must use False: it is the only path whose construction step can
be intercepted by our ambient meta context. With True, diffusers would hand
back fully materialized weights on every rank and the meta-shell design (shard
on meta -> to_empty only this rank's shard on GPU -> broadcast from rank0)
would silently degenerate into N full copies; the actor guards this with an
explicit "did not honor meta initialization" check.

The conflict: some model classes (e.g. WanTransformer3DModel) declare
_keep_in_fp32_modules -- submodules that must stay fp32 even when the rest
loads in bf16. diffusers implements that pinning only inside the True
path, and rather than silently dropping the pin it raises up-front:
ValueError: low_cpu_mem_usage cannot be False when keep_in_fp32_modules is True. So non-rank0 (forced False) cannot load such a class as-is.

How we handle it: on non-rank0 only, _keep_in_fp32_modules is temporarily
cleared on the class for the duration of the from_pretrained call. This is
safe because non-rank0 params are meta placeholders anyway -- authoritative
dtypes come from rank0 (which loads with True and honors the pin), and
sync_model_dtypes broadcasts rank0's per-tensor dtypes to every rank before
sharding, so all shards agree.